from fastapi import FastAPI, Request, Form
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse
import yt_dlp
import os
from pathlib import Path
import browser_cookie3  # 新增：用于获取浏览器 cookies

app = FastAPI()

# 静态文件和模板配置
app.mount("/static", StaticFiles(directory="static"), name="static")
app.mount("/downloads", StaticFiles(directory="downloads"), name="downloads")
templates = Jinja2Templates(directory="templates")

# 创建下载目录
Path("downloads/video").mkdir(parents=True, exist_ok=True)
Path("downloads/audio").mkdir(parents=True, exist_ok=True)

def get_cookies():
    """获取浏览器的 YouTube cookies"""
    try:
        # 尝试从多个浏览器获取 cookies
        browsers = [
            ('chrome', browser_cookie3.chrome),
            ('firefox', browser_cookie3.firefox),
            ('edge', browser_cookie3.edge),
            ('safari', browser_cookie3.safari),
        ]
        
        for browser_name, browser_func in browsers:
            try:
                cookies = browser_func(domain_name='.youtube.com')
                return {cookie.name: cookie.value for cookie in cookies}
            except:
                continue
                
        return {}
    except:
        return {}

def get_video_info(url: str):
    """获取视频信息"""
    ydl_opts = {
        'quiet': True,
        'no_warnings': True,
        'extract_flat': True,
        'cookiesfrombrowser': ('chrome',),  # 添加 cookies 支持
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        try:
            info = ydl.extract_info(url, download=False)
            return {
                'title': info.get('title', ''),
                'duration': info.get('duration', 0),
                'thumbnail': info.get('thumbnail', ''),
                'description': info.get('description', ''),
                'view_count': info.get('view_count', 0),
                'uploader': info.get('uploader', ''),
            }
        except Exception as e:
            return {'error': str(e)}

def download_video(url: str, download_type: str):
    """下载视频或音频"""
    video_info = get_video_info(url)
    if 'error' in video_info:
        return {'error': video_info['error']}

    safe_title = "".join(x for x in video_info['title'] if x.isalnum() or x in (' ', '-', '_'))[:50]
    
    # 基础下载选项
    ydl_opts = {
        'cookiesfrombrowser': ('chrome',),  # 添加 cookies 支持
        'quiet': False,
        'no_warnings': True,
    }
    
    if download_type == 'video':
        ydl_opts.update({
            'format': 'best',
            'outtmpl': f'downloads/video/{safe_title}.%(ext)s',
        })
    else:  # audio
        ydl_opts.update({
            'format': 'bestaudio/best',
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',
            }],
            'outtmpl': f'downloads/audio/{safe_title}.%(ext)s',
        })

    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([url])
            
        # 获取下载后的文件路径
        if download_type == 'video':
            file_path = f'/downloads/video/{safe_title}.mp4'
        else:
            file_path = f'/downloads/audio/{safe_title}.mp3'
            
        return {
            **video_info,
            'file_path': file_path,
            'download_type': download_type,
        }
    except Exception as e:
        return {'error': str(e)}

@app.get("/")
async def home(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

@app.post("/download")
async def download(url: str = Form(...), download_type: str = Form(...)):
    result = download_video(url, download_type)
    return JSONResponse(content=result) 